Skip to content

Fix memory leak in Python extension module initialization #91

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from

Conversation

ksmasoun
Copy link

@ksmasoun ksmasoun commented Aug 8, 2025

Summary

Fixes a memory leak in the Metal Python extension module initialization by properly handling reference counting for type objects.

Problem

The original code used PyModule_AddObject() which steals references, but the code was still decrementing reference counts in the error path, leading to potential memory leaks and incorrect reference counting.

Solution

  • Replace PyModule_AddObject() with PyModule_AddObjectRef() which handles reference counting correctly
  • Add proper Py_DECREF() calls after successful module addition since PyModule_AddObjectRef() increments references
  • Maintain existing error handling path with Py_XDECREF() calls

Changes

  • Updated PyInit__metal() function in gpt_oss/metal/python/module.c
  • Improved reference counting for Model, Tokenizer, and Context type objects
  • Added explanatory comments for clarity

Py_DECREF(model_type);
Py_DECREF(tokenizer_type);
Py_DECREF(context_type);

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Duplicating deallocation code is undesirable. Instead, rename error: label to cleanup: and make it execute on both success and failure, similarly to

cleanup:
if (fd != -1) {
close(fd);
fd = -1;
}
gptoss_model_release(model); // does nothing if model is NULL
gptoss_tokenizer_release(tokenizer); // does nothing if tokenizer is NULL
return status;
}

@ksmasoun
Copy link
Author

ksmasoun commented Aug 12, 2025

@Maratyszcza done!

@ksmasoun ksmasoun requested a review from Maratyszcza August 12, 2025 21:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants